#!/usr/bin/env python

import sys
import argparse
import os

parser = argparse.ArgumentParser(description='This script is for creating a standard iToL "label" and/or "branch" color file when given the IDs of the genomes you want to color. For version info, run `bit-version`.')

required = parser.add_argument_group('required arguments')

required.add_argument("-g", "--target-genomes", help='Single-column file with the genomes to color (need to match the IDs in the tree file, with no ">")', action="store", required=True)
parser.add_argument("-w", "--what-to-color", help='What to color, must be: "branches", "labels", or "both" (default: "both")', action="store", dest="to_color", default="both")
parser.add_argument("-c", "--color", help='Color to use, pre-baked options include "blue", "green", "red", "purple", or "black", or can provide the hexcode (default: "blue", of course, \'cause it\'s the best)', action="store", dest="color", default="blue")
parser.add_argument("-l", "--line-weight", help='Line weight if coloring branches (default: "2")', action="store", dest="line_weight", default="2")
parser.add_argument("-o", "--output_file", help='Output file for iToL (default: "iToL-colors.txt")', action="store", dest="output_file", default="iToL-colors.txt")

if len(sys.argv)==1:
    parser.print_help(sys.stderr)
    sys.exit(0)

args = parser.parse_args()

if args.color == "blue":
    col = "#434da7"
elif args.color == "green":
    col = "#48a743"
elif args.color == "red":
    col = "#c01820"
elif args.color == "purple":
    col = "#512f9c"
elif args.color == "black":
    col = "#000000"
else:
    if not args.color.startswith("#"):
        print("\n\tSorry, we're not prepared to handle \"" + str(args.color) + "\" as the color... :(\n")
        parser.print_help(sys.stderr)
        sys.exit(1)
    else:
        col = args.color

if args.to_color not in ["both", "branches", "labels"]:
    print("\n\tSorry, we're not prepared to handle \"" + str(args.to_color) + "\" as the argument for what to color... :(\n")
    parser.print_help(sys.stderr)
    sys.exit(1)

try:
    line_weight = float(args.line_weight)
except ValueError:
    print("\n\tSorry, " + str(args.line_weight) + " doesn't appear to be a number... :(\n")
    parser.print_help(sys.stderr)
    sys.exit(1)

target_list = []

with open(args.target_genomes, "r") as target_genomes:
    for genome in target_genomes:
        target_list.append(genome.strip())

out_file = open(args.output_file, "w")

out_file.write("TREE_COLORS\nSEPARATOR TAB\nDATA\n\n")

# writing lines for coloring labels if needed
if args.to_color in ["both", "labels"]:

    for target in target_list:
        out_file.write(str(target) + "\tlabel\t" + str(col) + "\tbold\n")

# writing lines for coloring branches if needed
if args.to_color in ["both", "branches"]:

    for target in target_list:
        out_file.write(str(target) + "\tbranch\t" + str(col) + "\tnormal\t" + str(line_weight) + "\n")

out_file.close()
